package com.huawei.spare.part.replace.rx;

import io.reactivex.Observable;
import io.reactivex.functions.Function;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class RetryWithDelay implements Function<Observable<? extends Throwable>, Observable<?>> {
  private static final Logger log = LoggerFactory.getLogger(RetryWithDelay.class);

  private final int maxRetries;
  private final int retryDelayMillis;
  private AtomicInteger retryCount;
  private final String logMessage;

  public RetryWithDelay(final int maxRetries, final int retryDelayMillis) {
    this(maxRetries, retryDelayMillis, "Task failed, retry times: %d, reason: %s");
  }

  public RetryWithDelay(final int maxRetries, final int retryDelayMillis, String logMessage) {
    this.maxRetries = maxRetries;
    this.retryDelayMillis = retryDelayMillis;
    this.retryCount = new AtomicInteger(0);
    this.logMessage = logMessage;
  }

  @Override
  public Observable<?> apply(final Observable<? extends Throwable> attempts) {
    return attempts
        .flatMap((Function<Throwable, Observable<?>>) throwable -> {
          if (retryCount.incrementAndGet() <= maxRetries) {
            // When this Observable calls onNext, the original
            // Observable will be retried (i.e. re-subscribed).
            log.info(String.format(this.logMessage, retryCount.get(), throwable.getMessage()));
            return Observable.timer(retryDelayMillis, TimeUnit.MILLISECONDS);
          }
          // Max retries hit. Just pass the error along.
          retryCount.set(0);  // reset retry counts
          return Observable.error(throwable);
        });
  }
}